#pragma once

#include <pcl/registration/icp.h>
#include "registration_interface.hpp"

namespace TESTICP
{
    class ICPRegistration : public RegistrationInterface
    {
    public:
        ICPRegistration(const YAML::Node& node);
        ICPRegistration(float max_corr_dist, float trans_eps, float euc_fitness_eps, int max_iter);

        void setInputTarget(const CloudPtr& target) override;
        void setInputSource(const CloudPtr& source) override;

        bool align( CloudType& output_ptr, const Eigen::Matrix4f& predict_pose = Eigen::Matrix4f::Identity()) override;

        inline Eigen::Matrix4f getFinalTransformation() { return icp_ptr_->getFinalTransformation(); }
        inline float getFitnessScore() { return icp_ptr_->getFitnessScore(); }

    private:
        bool SetRegistrationParam(float max_corr_dist, float trans_eps, float euc_fitness_eps, int max_iter);

    private:
        pcl::IterativeClosestPoint<PointType, PointType>::Ptr icp_ptr_;
    };

    ICPRegistration::ICPRegistration(
        const YAML::Node& node) : icp_ptr_(new pcl::IterativeClosestPoint<PointType, PointType>()) {

        float max_corr_dist = node["max_corr_dist"].as<float>();
        float trans_eps = node["trans_eps"].as<float>();
        float euc_fitness_eps = node["euc_fitness_eps"].as<float>();
        int max_iter = node["max_iter"].as<int>();

        SetRegistrationParam(max_corr_dist, trans_eps, euc_fitness_eps, max_iter);
    }

    ICPRegistration::ICPRegistration(float max_corr_dist, float trans_eps, float euc_fitness_eps, int max_iter)
        : icp_ptr_(new pcl::IterativeClosestPoint<PointType, PointType>()) {

        SetRegistrationParam(max_corr_dist, trans_eps, euc_fitness_eps, max_iter);
    }

    bool ICPRegistration::SetRegistrationParam( float max_corr_dist, float trans_eps, float euc_fitness_eps, int max_iter) {

        icp_ptr_->setMaxCorrespondenceDistance(max_corr_dist);
        icp_ptr_->setTransformationEpsilon(trans_eps);
        icp_ptr_->setEuclideanFitnessEpsilon(euc_fitness_eps);
        icp_ptr_->setMaximumIterations(max_iter);

        std::cout << "ICP params:" << std::endl
            << "max_corr_dist: " << max_corr_dist << ", "
            << "trans_eps: " << trans_eps << ", "
            << "euc_fitness_eps: " << euc_fitness_eps << ", "
            << "max_iter: " << max_iter
            << std::endl;

        return true;
    }

    void ICPRegistration::setInputTarget(const CloudPtr& target) {
        icp_ptr_->setInputTarget(target);
    }

    void ICPRegistration::setInputSource(const CloudPtr& source) {
        icp_ptr_->setInputSource(source);
    }

    bool ICPRegistration::align( CloudType& output_ptr, const Eigen::Matrix4f& predict_pose) {
        icp_ptr_->align(output_ptr, predict_pose);
        return true;
    }

}
